package codingforgreat.class18;

import java.util.HashSet;
import java.util.PriorityQueue;

public class Class04_TopKSumCrossTwoArrays {
    //两个有序数组的topk问题
    //思路：dp[i][j] 最角落放进堆中 每弹出一个把相邻的左边和上边放入堆中，用set去重
    public class Node{
        int index1;
        int index2;
        int sum;

        public Node(int index1, int index2, int sum) {
            this.index1 = index1;
            this.index2 = index2;
            this.sum = sum;
        }
    }
    public  int[] topKSum(int[] arr1, int[] arr2, int topK) {
        if (arr1 == null || arr2 == null || topK < 1) {
            return null;
        }
        int N = arr1.length;
        int M = arr2.length;
        topK = Math.min(topK, N * M);
        int[] res = new int[topK];
        int resIndex = 0;
        HashSet<String> set = new HashSet();
        PriorityQueue<Node> queue = new PriorityQueue<>((x,y)->{
            return y.sum - x.sum;
        });
        int i1 = N - 1;
        int i2 = M - 1;
        queue.add(new Node(i1, i2, arr1[i1] + arr2[i2]));
        set.add(i1 + "_" + i2);
        while(resIndex != topK){
            Node curNode = queue.poll();
            res[resIndex++] = curNode.sum;
            i1 = curNode.index1;
            i2 = curNode.index2;
            set.remove(i1 + "_" + i2);
            if (i1 - 1 >= 0 && !set.contains((i1 - 1) + "_" + i2)) {
                set.add((i1 - 1) + "_" + i2);
                queue.add(new Node(i1 - 1, i2, arr1[i1 - 1] + arr2[i2]));
            }
            if (i2 - 1 >= 0 && !set.contains(i1 + "_" + (i2 - 1))) {
                set.add(i1 + "_" + (i2 - 1));
                queue.add(new Node(i1, i2 - 1, arr1[i1] + arr2[i2 - 1]));
            }
        }
        return res;
    }
}
